{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8d70339c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.stats import wasserstein_distance\n",
    "import helpers as ph\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import statsmodels.api as sm\n",
    "import dataframe_image as dfi\n",
    "\n",
    "sns.set_style('darkgrid')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "997f5650",
   "metadata": {},
   "outputs": [],
   "source": [
    "RESULTS_DIR = f'./data/distributions/'\n",
    "CONTEXT = 'default'\n",
    "SAVEFIG = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4f60044",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_df = []\n",
    "for context in ['default', 'steer-qa', 'steer-bio', 'steer-portray']:\n",
    "    \n",
    "    cdf = pd.read_csv(os.path.join(RESULTS_DIR, f'Pew_American_Trends_Panel_disagreement_500_{context}_combined.csv'))\n",
    "    cdf['survey'] = 'Pew_American_Trends_Panel_disagreement_500'\n",
    "    cdf['context'] = context\n",
    "    if context != 'default':\n",
    "        cdf = cdf[cdf['group'] == cdf['steer_group']]\n",
    "    combined_df.append(cdf)\n",
    "combined_df = pd.concat(combined_df)\n",
    "combined_df['Source'] = combined_df.apply(lambda x: 'AI21 Labs' if 'j1-' in x['model_name'].lower() else 'OpenAI',\n",
    "                                          axis=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91894898",
   "metadata": {},
   "source": [
    "## Measure steerability"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08f1e19f",
   "metadata": {},
   "outputs": [],
   "source": [
    "KEYS = ['model_name', 'context', 'attribute', 'group', 'group_order', 'model_order']\n",
    "\n",
    "grouped = combined_df.groupby(KEYS, as_index=False).agg({'WD': np.mean}) \\\n",
    "         .sort_values(by=['context', 'model_order', 'group_order'])\n",
    "grouped['Rep'] = 1 - grouped['WD']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1d921b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "unsteered = grouped[grouped['context'] == 'default'].rename(columns={'WD': 'WD_d',\n",
    "                                                                     'Rep': 'Rep_d'})\n",
    "steered = grouped[grouped['context'] != 'default'].sort_values(by='Rep')\n",
    "steered = steered.groupby([k for k in KEYS if k != 'context'], as_index=False).last()\\\n",
    "                 .rename(columns={'WD': 'WD_s', 'Rep': 'Rep_s'})\n",
    "result = pd.merge(unsteered, steered, on=[k for k in KEYS if k != 'context']) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5392a246",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(6, 4))\n",
    "sns.set_style('whitegrid')\n",
    "\n",
    "for model in ph.MODEL_NAMES.values():\n",
    "    if model not in set(result['model_name'].values): continue \n",
    "    c = result[result['model_name'] == model]\n",
    "    \n",
    "    reg = sm.OLS(c['Rep_s'], c['Rep_d'])\n",
    "    slope = reg.fit().params['Rep_d']\n",
    "\n",
    "    sns.regplot(data=c, x='Rep_d', y='Rep_s', ax=ax, \n",
    "                label=model, \n",
    "                line_kws={'linewidth': 3},\n",
    "                scatter_kws={'s': 14})\n",
    "    \n",
    "xx = np.linspace(0.66, 0.9, 10)\n",
    "plt.legend(loc=4, fontsize=9, ncol=2)\n",
    "ax.plot(xx, xx, 'k--')\n",
    "ax.set_xlim([0.68, 0.88])\n",
    "ax.set_ylim([0.68, 0.88])\n",
    "\n",
    "plt.xlabel('Default subgroup representativeness', fontsize=12)\n",
    "plt.ylabel('Steered subgroup representativeness', fontsize=12)\n",
    "plt.grid(linestyle='-', linewidth=0.5)\n",
    "if SAVEFIG: plt.savefig('./figures/steerability.png', bbox_inches=\"tight\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ddc97eb",
   "metadata": {},
   "source": [
    "## Steerability by topic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "380fd3a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "topic_info = np.load('./data/human_resp/topic_mapping.npy', allow_pickle=True).item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a6ba501d",
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_df['topic'] = combined_df.apply(lambda x: topic_info[x['question']]['cg'], axis=1)\n",
    "combined_df = combined_df.explode(['topic'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c01ff484",
   "metadata": {},
   "outputs": [],
   "source": [
    "KEYS = ['Source', 'model_name', 'context', 'attribute', 'group', 'group_order', 'model_order', 'topic']\n",
    "\n",
    "grouped_topic = combined_df.groupby(KEYS, as_index=False).agg({'WD': np.mean}) \\\n",
    "         .sort_values(by=['context', 'model_order', 'group_order'])\n",
    "grouped_topic['Rep'] = 1 - grouped_topic['WD']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da97b819",
   "metadata": {},
   "outputs": [],
   "source": [
    "steered_topic = grouped_topic[grouped_topic['context'] != 'default'].sort_values(by='Rep')\n",
    "steered_topic = steered_topic.groupby([k for k in KEYS if k != 'context'], as_index=False).last()\\\n",
    "                 .rename(columns={'WD': 'WD_s', 'Rep': 'Rep_s', 'model_name': ''})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95bec6da",
   "metadata": {},
   "outputs": [],
   "source": [
    "steered_topic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f1f1469",
   "metadata": {},
   "outputs": [],
   "source": [
    "styles = ph.VIS_STYLES\n",
    "\n",
    "styles[-1]['props'][-1] = (styles[-1]['props'][-1][0], \"105%\")\n",
    "for attribute in np.unique(steered_topic['attribute']):\n",
    "    print(attribute)\n",
    "    table = pd.pivot_table(steered_topic[steered_topic['attribute'] == attribute], \n",
    "                           columns=['Source', ''], \n",
    "                       index='group',\n",
    "                       values='Rep_s', \n",
    "                       sort=False)\n",
    "    table_vis = table.style.background_gradient('Reds', axis=1).set_table_styles(styles)  \\\n",
    "                           .set_properties(**{\"font-size\":\"0.8rem\"}).format(precision=3)\n",
    "    if SAVEFIG: dfi.export(table_vis, f'./figures/steerability_{attribute}.png')\n",
    "\n",
    "    display(table_vis)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a4242e7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
